# environment.py

import numpy as np
from config import NUM_TARGETS, FREQ_RANGE, MAX_JAM_BANDWIDTH

class JammingEnvironment:
    def __init__(self):
        # 初始化目标信息（威胁系数、频率分布）
        self.targets = self._initialize_targets()
        self.state = self._get_initial_state()

    def _initialize_targets(self):
        # 模拟6个跳频目标的频率分布和威胁系数
        np.random.seed(42)  # 固定随机种子
        freq_centers = np.linspace(*FREQ_RANGE, NUM_TARGETS)
        freq_spread = np.random.randint(1, 5, NUM_TARGETS)  # 频点分布范围
        return [
            {"threat": threat, "freq_center": center, "spread": spread}
            for threat, center, spread in zip(
                range(NUM_TARGETS, 0, -1),
                freq_centers,
                freq_spread
            )
        ]

    def _get_initial_state(self):
        # 构造状态空间：威胁系数 + 频段划分信息
        freq_segments = self._divide_frequency_segments()
        return {
            "threats": [t["threat"] for t in self.targets],
            "segments": freq_segments,
        }

    def _divide_frequency_segments(self):
        # 按照频率分布划分子频段
        freqs = [t["freq_center"] for t in self.targets]
        freqs.sort()
        segments = []
        current_segment = [freqs[0]]
        for i in range(1, len(freqs)):
            if freqs[i] - freqs[i-1] <= MAX_JAM_BANDWIDTH:
                current_segment.append(freqs[i])
            else:
                segments.append(current_segment)
                current_segment = [freqs[i]]
        segments.append(current_segment)
        return segments

    def step(self, action_freq, action_band):
        # 执行干扰动作，返回奖励和新状态
        # 这里简化逻辑，实际需结合公式(1)-(5)计算干信比
        reward = self._calculate_reward(action_freq, action_band)
        next_state = self._get_next_state()
        done = self._check_done()
        return next_state, reward, done

    def _calculate_reward(self, action_freq, action_band):
        # 根据干扰效果计算奖励（参考论文公式）
        # 示例：成功干扰目标获得正奖励，否则负奖励
        return np.random.choice([10, -2])  # 模拟奖励值

    def _get_next_state(self):
        # 更新状态（模拟干扰后的频段变化）
        return self.state

    def _check_done(self):
        # 判断是否完成干扰任务
        return np.random.rand() > 0.1  # 模拟完成条件

    def reset(self):
        # 重置环境
        self.state = self._get_initial_state()
        return self.state